//This file contains a smart contract: ModelContract.
//It handles model download, model aggression ,model upload and model query business.
package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/hyperledger/fabric-contract-api-go/contractapi"
	"strconv"
	"time"
)

//ModelContract is an instantiation of smart contract.
type ModelContract struct {
	contractapi.Contract
}

//ModelInformation will be used as the key in the key-value pair of a local model transaction.
//And the value in the key-value pair of a local model transaction is TrainingParameters.
type ModelInformation struct {
	//if true, this is the key of a local model.If false, this is the key of a global model.
	IsLocalModel bool
	Iteration int
	Producer string
	CreatedTime string
}

//GlobalModel will be used as the value in the key-value pair of a global model transaction.
//And the key in the key-value of a global model transaction is the string "Iteration".
//Thus, if a peer need to query information about latest global model or Iter, it will query key "Iteration" in the ledger.
type GlobalModel struct {
	ModelInformation
	//TPAsString is the result of type TrainingParameters being marshalled and converted to string type.
	TPAsString string
}

//Iter is an abbreviation of word "iteration". It is a global variable showing the current round of federated learning.
//As is known, "querying the global model", "training the model locally", "submitting the local model", and "aggregating a new global model" are four steps of a round of federated learning.
//And This process will continue for multiple rounds.
//Thus, Iter tells us the current round, or in other words, the current iteration.
//In case of single point of failure and nodes out of synchronization, function IterUpdater is called to keep Iter up-to-date whenever a node use ModelContract to query or submit.
var Iter =0


/***Following are methods of ModelContract.They can be divided into methods of submitting and methods of querying.***/

/*1.methods of submitting*/

//SubmitLocalModel submit a local model into ledger
func (mc *ModelContract) SubmitLocalModel(ctx contractapi.TransactionContextInterface,producer string,value string) (string,error) {
	//update Iter
	err:= IterUpdater(ctx)
	if err != nil {
		return "",fmt.Errorf("SubmitLocalModel failed for function IterUpdater failed: %s",err.Error())
	}

	//create key for local model transaction
	NowTime:=time.Now()
	mn:= ModelInformation{true,Iter,producer,NowTime.Format("2006-01-02 15:04:05")}
	ComKey,err:=ctx.GetStub().CreateCompositeKey("ModelInformation",[]string{strconv.FormatBool(mn.IsLocalModel),strconv.Itoa(mn.Iteration),mn.Producer,mn.CreatedTime})
	if err!=nil {
		return "",fmt.Errorf("cannot create composite key: %s",err.Error())
	}

	//submit local model into ledger
	err = ctx.GetStub().PutState(ComKey, []byte(value))
	if err != nil {
		return "",fmt.Errorf("unable to interact with world state :%s",err.Error())
	}

	//inform user
	Info:=fmt.Sprintf("successfully put a local model to the ledger:maker %s,Iteration %d,created time %s",producer,Iter,NowTime.Format("2006-01-02 15:04:05"))
	return Info,nil
}

//SubmitGlobalModel submit a global model into ledger.
//It is worth noting that, SubmitGlobalModel allows nodes to submit any user-specified global models.
//Thus, SubmitGlobalModel method can only be called before 1st round, to submit original global model for starting federated learning process.
//Subsequent global models should only be created by calling AggregationAndSubmit method,
//because AggregationAndSubmit will only submit certain global models that contains result of aggregation, instead of user-specified.
func (mc *ModelContract) SubmitGlobalModel(ctx contractapi.TransactionContextInterface,producer string,value string) (string,error) {
	//update Iter
	err:= IterUpdater(ctx)

	//if this is not 1st round, then calling of SubmitGlobalModel is forbidden.
	if Iter != 0 {
		return "", errors.New("FAILED!!! Because once the learning process has been started, you are not allowed to submit global models that you specified. Please call AggregationAndSubmit method to submit a global model")
	}

	if err != nil {
		return "",fmt.Errorf("SubmitGlobalModel failed for function IterUpdater failed: %s",err.Error())
	}

	//call GlobalModelSubmitter to submit global model
	Info,err:= GlobalModelSubmitter(ctx,Iter,producer,value)
	if err != nil {
		return "",fmt.Errorf("function GlobalModelSubmitter error: %s",err)
	}
	return Info,nil
}

//AggregationAndSubmit executes the following three steps:
//First, it queries all new local models that was created after latest global model.
//Second, it aggregates these new local models into a global model.
//Third,it submit the global model into ledger.
func (mc *ModelContract) AggregationAndSubmit(ctx contractapi.TransactionContextInterface,producer string) (string,error) {
	//update Iter
	err:= IterUpdater(ctx)
	if err != nil {
		return "",fmt.Errorf("AggregationAndSubmit failed because function IterUpdater failed: %s",err.Error())
	}

	//aggregate TrainingParameters of local models
	AverData,err:= Aggregator(ctx,strconv.Itoa(Iter))
	if err!=nil {
		return "",fmt.Errorf("AggregationAndSubmit failed because function Aggregator failed: %s",err.Error())
	}
	if AverData == nil {
		return "There is no new local model in the ledger",nil
	}

	//turn type TrainingParameters into json
	AverDataAsBytes,err:=json.Marshal(AverData)
	if err!=nil {
		return "",fmt.Errorf("failed to marshal: %s",err.Error())
	}

	//call GlobalModelSubmitter to submit global model
	Info,err:= GlobalModelSubmitter(ctx,Iter,producer,string(AverDataAsBytes))
	if err != nil {
		return "",fmt.Errorf("function GlobalModelSubmitter error: %s",err)
	}

	//return result
	return Info,nil
}

/*2.methods of querying*/

//QueryLocalModelSinceLatestGlobalModel queries all new local models that was created after latest global model
func (mc *ModelContract) QueryLocalModelSinceLatestGlobalModel(ctx contractapi.TransactionContextInterface) ([]string,error) {
	//update Iter
	err:= IterUpdater(ctx)
	if err != nil {
		return nil,fmt.Errorf("QueryLocalModelSinceLatestGlobalModel failed for function IterUpdater failed: %s",err.Error())
	}

	//query new local models in ledger
	KeyValuePairs,err:= LocalModelQueryer(ctx,strconv.Itoa(Iter))
	if err != nil {
		return nil,fmt.Errorf("function LocalModelQueryer returns an error: %s",err.Error())
	}

	//return result
	var results []string
	for key,value:=range KeyValuePairs{
		result:="{KEY: "+key+" && VALUE: "+string(value)+"}"
		results = append(results, result)
	}
	return results, nil
}

//QueryLocalModelByIteration queries all local models of given iteration
func (mc *ModelContract) QueryLocalModelByIteration(ctx contractapi.TransactionContextInterface, iteration string) ([]string, error) {
	//query local models in ledger of given iteration
	KeyValuePairs,err:= LocalModelQueryer(ctx,iteration)
	if err != nil {
		return nil,fmt.Errorf("function LocalModelQueryer returns an error: %s",err.Error())
	}

	//return result
	var results []string
	for key,value:=range KeyValuePairs{
		result:="{KEY: "+key+" && VALUE: "+string(value)+"}"
		results = append(results, result)
	}
	return results, nil
}

//QueryLatestGlobalModel queries latest global model
func (mc *ModelContract) QueryLatestGlobalModel(ctx contractapi.TransactionContextInterface) (string,error) {
	//update Iter
	err:= IterUpdater(ctx)
	if err != nil {
		return "",fmt.Errorf("QueryLatestGlobalModel failed for function IterUpdater failed: %s",err.Error())
	}

	//query latest global model in ledger
	queryResponse,err:=ctx.GetStub().GetState("Iteration")
	if err != nil {
		return "",fmt.Errorf("GetState failed: %s",err.Error())
	}
	if queryResponse == nil {
		return "There is no global model in the ledger",nil
	}

	//unmarshal
	GM:=new(GlobalModel)
	err=json.Unmarshal(queryResponse,GM)
	if err != nil {
		return "", fmt.Errorf("cannot unmarshal value of global model query response: %s",err.Error())
	}

	//turn field ModelInformation of GM(GM's type is GlobalModel) into string type
	ComKey,err:=ctx.GetStub().CreateCompositeKey("ModelInformation",[]string{strconv.FormatBool(GM.IsLocalModel),strconv.Itoa(GM.Iteration),GM.Producer,GM.CreatedTime})
	if err != nil {
		return "", fmt.Errorf("cannot create composite key: %s",err.Error())
	}

	//return result
	result:="{KEY: "+ComKey+" VALUE: "+GM.TPAsString+"}"
	return result,nil
}

//QueryAllModel queries all models in the ledger
func (mc *ModelContract) QueryAllModel(ctx contractapi.TransactionContextInterface) ([]string,error) {
	//query all local models in ledger
	LResultsIterator,err:=ctx.GetStub().GetStateByPartialCompositeKey("ModelInformation",[]string{})
	if err != nil {
		return nil, fmt.Errorf("cannot read all local models because of searching failed: %s",err.Error())
	}
	defer func() {
		_=LResultsIterator.Close()
	}()

	//query all history of global model in ledger
	GResultsIterator,err:=ctx.GetStub().GetHistoryForKey("Iteration")
	if err != nil {
		return nil, fmt.Errorf("cannot read all global models because of searching failed: %s",err.Error())
	}
	defer func() {
		_=GResultsIterator.Close()
	}()

	//new results
	var results []string

	//put all history of global model into results
	for GResultsIterator.HasNext() {
		queryResponse, err := GResultsIterator.Next()
		if err != nil {
			return nil, fmt.Errorf("function Next of GResultsIterator error: %s",err.Error())
		}
		GM:=new(GlobalModel)
		err=json.Unmarshal(queryResponse.Value,GM)
		if err != nil {
			return nil, fmt.Errorf("cannot unmarshal value of global model query response: %s",err.Error())
		}
		ComKey,err:=ctx.GetStub().CreateCompositeKey("ModelInformation",[]string{strconv.FormatBool(GM.IsLocalModel),strconv.Itoa(GM.Iteration),GM.Producer,GM.CreatedTime})
		if err != nil {
			return nil, fmt.Errorf("cannot create composite key: %s",err.Error())
		}
		result:="{KEY: "+ComKey+" && VALUE: "+GM.TPAsString+"}"
		results = append(results, result)
	}

	//put all local models into results
	for LResultsIterator.HasNext() {
		queryResponse, err := LResultsIterator.Next()
		if err != nil {
			return nil, fmt.Errorf("function Next of LResultsIterator error: %s",err.Error())
		}

		result:="{KEY: "+queryResponse.Key+" && VALUE: "+string(queryResponse.Value)+"}"
		results = append(results, result)
	}

	//return results
	return results, nil
}

//QueryAmountOfLocalModelSinceLatestGlobalModel queries number of new local models that was created after latest global model
func (mc *ModelContract) QueryAmountOfLocalModelSinceLatestGlobalModel(ctx contractapi.TransactionContextInterface) (string,error) {
	//update Iter
	err:= IterUpdater(ctx)
	if err != nil {
		return "",fmt.Errorf("QueryAmountOfLocalModelSinceLatestGlobalModel failed for function IterUpdater failed: %s",err.Error())
	}

	//query new local models in ledger
	resultsIterator,err:=ctx.GetStub().GetStateByPartialCompositeKey("ModelInformation",[]string{strconv.FormatBool(true),strconv.Itoa(Iter)})
	if err != nil {
		return "", fmt.Errorf("cannot read local models because of searching failed: %s",err.Error())
	}
	defer func() {
		_=resultsIterator.Close()
	}()

	//count
	lens:=0
	for resultsIterator.HasNext() {
		_, err := resultsIterator.Next()
		lens++
		if err != nil {
			return "", fmt.Errorf("error while traverse resultsIterator: %s",err.Error())
		}
	}

	//return result
	return strconv.Itoa(lens), nil
}

//QueryIteration queries current iteration.
func (mc *ModelContract) QueryIteration (ctx contractapi.TransactionContextInterface) (string,error) {
	//update Iter
	err:= IterUpdater(ctx)
	if err != nil {
		return "",fmt.Errorf("QueryIteration failed for function IterUpdater failed: %s",err.Error())
	}

	//return result
	return strconv.Itoa(Iter),nil
}

//StimulationOfAggregation is similar to AggregationAndSubmit.
//However, it only return the result of aggregation to user's terminal, instead of submitting it to ledger.
//Which is to say, it won't execute the third step of AggregationAndSubmit.
func (mc *ModelContract) StimulationOfAggregation(ctx contractapi.TransactionContextInterface) (string,error) {
	//update Iter
	err:= IterUpdater(ctx)
	if err != nil {
		return "",fmt.Errorf("StimulationOfAggregation failed for function IterUpdater failed: %s",err.Error())
	}

	//aggregate TrainingParameters of local models
	AverData,err:= Aggregator(ctx,strconv.Itoa(Iter))
	if err!=nil {
		return "",fmt.Errorf("function Aggregator error: %s",err.Error())
	}
	if AverData == nil {
		return "There is no new local model in the ledger",nil
	}

	//turn type TrainingParameters into json
	AverDataAsBytes,err:=json.Marshal(AverData)
	if err!=nil {
		return "",fmt.Errorf("failed to marshal: %s",err.Error())
	}

	//return result of aggregation stimulation to terminal
	return string(AverDataAsBytes),nil
}
